#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
动作学习管理器 - Action Learning Manager

该模块提供六足机器人的动作学习功能，包括：
- 动作记录和保存
- 动作序列播放
- 动作文件管理
- 学习模式控制

Author: YAHBOOM | Gentle Xu
Date: 2025/09
"""

import os
import time
import json
from typing import List, Tuple, Optional, Dict, Any
from datetime import datetime


class ActionLearningManager:
    """动作学习管理器类"""
    
    def __init__(self, robot_controller, data_dir: str = "action_data"):
        """
        初始化动作学习管理器
        
        Args:
            robot_controller: 机器人控制器实例
            data_dir: 动作数据存储目录
        """
        self.robot_controller = robot_controller
        self.data_dir = data_dir
        self.is_learning = False
        self.current_action_sequence = []
        
        # 创建数据目录
        if not os.path.exists(self.data_dir):
            os.makedirs(self.data_dir)
        
        # 前腿舵机ID映射 (左前腿和右前腿)
        self.front_leg_motors = {
            'left_front': [1, 2, 3],    # 左前腿舵机ID
            'right_front': [16, 17, 18] # 右前腿舵机ID
        }
        
        # 前腿ID
        self.front_leg_ids = [1, 6]  # 左前腿ID=1, 右前腿ID=6
    
    def prepare_learning_posture(self) -> bool:
        """
        进入准备学习的姿态
        
        Returns:
            bool: 是否成功进入学习姿态
        """
        try:
            print("正在进入学习准备姿态...")
            
            # 重置机器人
            self.robot_controller.robot.reset()
            time.sleep(1)
            
            # 调整特定舵机到学习位置
            self.robot_controller.robot.motor(5, -60, 500)
            self.robot_controller.robot.motor(14, -60, 500)
            time.sleep(0.5)
            
            self.robot_controller.robot.motor(4, 30, 500)
            self.robot_controller.robot.motor(13, -30, 500)
            time.sleep(0.5)
            
            self.robot_controller.robot.motor(5, -35, 500)
            self.robot_controller.robot.motor(14, -35, 500)
            time.sleep(0.5)
            
            print("学习准备姿态设置完成")
            return True
            
        except Exception as e:
            print(f"设置学习姿态失败: {e}")
            return False
    
    def start_learning(self) -> bool:
        """
        开始学习模式，卸载前腿舵机
        
        Returns:
            bool: 是否成功开始学习
        """
        try:
            if self.is_learning:
                print("已经处于学习模式")
                return True
            
            print("开始动作学习模式...")
            
            # 卸载前腿舵机，使其可以手动调节
            for leg_id in self.front_leg_ids:
                self.robot_controller.robot.unload_leg(leg_id)
                time.sleep(0.1)
            
            self.is_learning = True
            self.current_action_sequence = []
            
            print("学习模式已启动，可以手动调节前腿姿态")
            return True
            
        except Exception as e:
            print(f"启动学习模式失败: {e}")
            return False
    
    def stop_learning(self) -> bool:
        """
        停止学习模式，重新加载前腿舵机
        
        Returns:
            bool: 是否成功停止学习
        """
        try:
            if not self.is_learning:
                print("当前不在学习模式")
                return True
            
            print("停止动作学习模式...")
            
            # 重新加载前腿舵机
            for leg_id in self.front_leg_ids:
                self.robot_controller.robot.load_leg(leg_id)
                time.sleep(0.1)
            
            self.is_learning = False
            
            print("学习模式已停止")
            return True
            
        except Exception as e:
            print(f"停止学习模式失败: {e}")
            return False
    
    def record_current_action(self) -> bool:
        """
        记录当前动作（前腿舵机角度）
        
        Returns:
            bool: 是否成功记录动作
        """
        try:
            if not self.is_learning:
                print("请先启动学习模式")
                return False
            
            # 直接从机器人读取所有舵机角度
            if not hasattr(self.robot_controller.robot, 'read_motor'):
                print("机器人不支持舵机角度读取")
                return False
            
            all_angles = self.robot_controller.robot.read_motor()
            if not all_angles or len(all_angles) < 18:
                print("无法读取舵机角度")
                return False
            
            # 提取前腿舵机角度 (6个舵机)
            action_angles = [0] * 6
            
            # 左前腿 (舵机ID: 1, 2, 3) -> 索引 0, 1, 2
            for i in range(3):
                servo_index = i  # 舵机1-3对应索引0-2
                action_angles[i] = all_angles[servo_index]
            
            # 右前腿 (舵机ID: 16, 17, 18) -> 索引 15, 16, 17
            for i in range(3):
                servo_index = 15 + i  # 舵机16-18对应索引15-17
                action_angles[3 + i] = all_angles[servo_index]
            
            # 添加到当前动作序列
            action_data = {
                'timestamp': datetime.now().isoformat(),
                'angles': action_angles,
                'description': f"Action {len(self.current_action_sequence) + 1}"
            }
            
            self.current_action_sequence.append(action_data)
            
            print(f"已记录动作 {len(self.current_action_sequence)}: {action_angles}")
            return True
            
        except Exception as e:
            print(f"记录动作失败: {e}")
            return False
    
    def save_action_sequence(self, filename: str) -> bool:
        """
        保存当前动作序列到文件
        
        Args:
            filename: 保存的文件名
            
        Returns:
            bool: 是否成功保存
        """
        try:
            if not self.current_action_sequence:
                print("没有动作序列可保存")
                return False
            
            filepath = os.path.join(self.data_dir, f"{filename}.json")
            
            sequence_data = {
                'metadata': {
                    'created_at': datetime.now().isoformat(),
                    'action_count': len(self.current_action_sequence),
                    'description': f"Action sequence with {len(self.current_action_sequence)} actions"
                },
                'actions': self.current_action_sequence
            }
            
            with open(filepath, 'w', encoding='utf-8') as f:
                json.dump(sequence_data, f, indent=2, ensure_ascii=False)
            
            print(f"动作序列已保存到: {filepath}")
            return True
            
        except Exception as e:
            print(f"保存动作序列失败: {e}")
            return False
    
    def load_action_sequence(self, filename: str) -> Optional[List[Dict]]:
        """
        从文件加载动作序列
        
        Args:
            filename: 文件名
            
        Returns:
            Optional[List[Dict]]: 动作序列数据，失败返回None
        """
        try:
            filepath = os.path.join(self.data_dir, f"{filename}.json")
            
            if not os.path.exists(filepath):
                print(f"文件不存在: {filepath}")
                return None
            
            with open(filepath, 'r', encoding='utf-8') as f:
                sequence_data = json.load(f)
            
            actions = sequence_data.get('actions', [])
            print(f"已加载动作序列: {len(actions)} 个动作")
            
            return actions
            
        except Exception as e:
            print(f"加载动作序列失败: {e}")
            return None
    
    def play_action_sequence(self, actions: List[Dict], speed: int = 800, delay: float = 1.0) -> bool:
        """
        播放动作序列
        
        Args:
            actions: 动作序列
            speed: 舵机移动速度
            delay: 动作间延迟时间(秒)
            
        Returns:
            bool: 是否成功播放
        """
        try:
            if not actions:
                print("没有动作可播放")
                return False
            
            if self.is_learning:
                print("请先停止学习模式")
                return False
            
            print(f"开始播放动作序列，共 {len(actions)} 个动作")
            
            # 前腿舵机ID映射
            motor_ids = [1, 2, 3, 16, 17, 18]
            
            for i, action in enumerate(actions):
                print(f"执行动作 {i+1}/{len(actions)}")
                
                angles = action.get('angles', [])
                if len(angles) != 6:
                    print(f"动作 {i+1} 数据格式错误，跳过")
                    continue
                
                # 同时控制6个前腿舵机
                for j, angle in enumerate(angles):
                    self.robot_controller.robot.motor(motor_ids[j], angle, speed)
                
                # 等待动作完成
                time.sleep(delay)
            
            print("动作序列播放完成")
            return True
            
        except Exception as e:
            print(f"播放动作序列失败: {e}")
            return False
    
    def clear_current_sequence(self) -> bool:
        """
        清空当前动作序列
        
        Returns:
            bool: 是否成功清空
        """
        try:
            self.current_action_sequence = []
            print("当前动作序列已清空")
            return True
        except Exception as e:
            print(f"清空动作序列失败: {e}")
            return False
    
    def delete_action_file(self, filename: str) -> bool:
        """
        删除动作文件
        
        Args:
            filename: 文件名
            
        Returns:
            bool: 是否成功删除
        """
        try:
            filepath = os.path.join(self.data_dir, f"{filename}.json")
            
            if os.path.exists(filepath):
                os.remove(filepath)
                print(f"已删除文件: {filepath}")
                return True
            else:
                print(f"文件不存在: {filepath}")
                return False
                
        except Exception as e:
            print(f"删除文件失败: {e}")
            return False
    
    def list_action_files(self) -> List[str]:
        """
        列出所有动作文件
        
        Returns:
            List[str]: 动作文件名列表
        """
        try:
            if not os.path.exists(self.data_dir):
                return []
            
            files = []
            for filename in os.listdir(self.data_dir):
                if filename.endswith('.json'):
                    files.append(filename[:-5])  # 去掉.json后缀
            
            return files
            
        except Exception as e:
            print(f"列出文件失败: {e}")
            return []
    
    def get_learning_status(self) -> Dict[str, Any]:
        """
        获取学习状态信息
        
        Returns:
            Dict[str, Any]: 状态信息
        """
        return {
            'is_learning': self.is_learning,
            'current_sequence_length': len(self.current_action_sequence),
            'data_directory': self.data_dir,
            'available_files': self.list_action_files()
        }
    
    def print_status(self):
        """
        打印当前状态信息
        """
        status = self.get_learning_status()
        print("\n=== 动作学习状态 ===")
        print(f"学习模式: {'开启' if status['is_learning'] else '关闭'}")
        print(f"当前序列长度: {status['current_sequence_length']}")
        print(f"数据目录: {status['data_directory']}")
        print(f"可用文件: {', '.join(status['available_files']) if status['available_files'] else '无'}")
        print("=====================\n")